import os.path
import os
os.path.join('./')
from data.base_dataset import BaseDataset, get_params, get_transform, normalize
from data.image_folder import make_dataset
from PIL import Image
import torch
import torch.nn.functional as F
import numpy as np
class AlignedDataset(BaseDataset):
    def initialize(self, opt):
        self.opt = opt
        self.root = opt.dataroot
        self.dir_A='/opt/data/private/yudong/dataset/MSRS_Fusion_source/VVVV/VIS'
        self.dir_M='/opt/data/private/yudong/dataset/MSRS_Fusion_source/VVVV_MASK/MASK'
        self.A_paths = sorted(make_dataset(self.dir_A))
        self.M_paths = sorted(make_dataset(self.dir_M))

        if opt.isTrain or opt.use_encoded_image:
            self.dir_B='/opt/data/private/yudong/dataset/MSRS_Fusion_source/IIII/IR'
            self.B_paths = sorted(make_dataset(self.dir_B))

        ### instance maps
        if not opt.no_instance:
            self.dir_inst = os.path.join(opt.dataroot, opt.phase + '_inst')
            self.inst_paths = sorted(make_dataset(self.dir_inst))

        ### load precomputed instance-wise encoded features
        if opt.load_features:                              
            self.dir_feat = os.path.join(opt.dataroot, opt.phase + '_feat')
            print('----------- loading features from %s ----------' % self.dir_feat)
            self.feat_paths = sorted(make_dataset(self.dir_feat))

        self.dataset_size = len(self.A_paths) 


    def get_detail(self, image):
        sobel_kernel_x = torch.tensor([[1, 0, -1],
                                       [2, 0, -2],
                                       [1, 0, -1]], dtype=torch.float32)

        sobel_kernel_y = torch.tensor([[1, 2, 1],
                                       [0, 0, 0],
                                       [-1, -2, -1]], dtype=torch.float32)

        sobel_kernel_x = sobel_kernel_x.view(1, 1, 3, 3)
        sobel_kernel_y = sobel_kernel_y.view(1, 1, 3, 3)
        image = torch.tensor(np.array(image), dtype=torch.float32)
        image = torch.unsqueeze(image, 0)
        image = torch.unsqueeze(image, 0)


        grad_x = F.conv2d(image, sobel_kernel_x, padding=1)
        grad_y = F.conv2d(image, sobel_kernel_y, padding=1)
        grad_magnitude = torch.sqrt(grad_x ** 2 + grad_y ** 2)
        grad_magnitude = torch.squeeze(grad_magnitude, 0)
        return Image.fromarray(grad_magnitude.numpy()[0])

    def __getitem__(self, index):
        A_path = self.A_paths[index]              
        A = Image.open(A_path)

        A_gray=Image.open(A_path).convert('L')
        detail = self.get_detail(A_gray)

        params = get_params(self.opt, A.size)
        if self.opt.label_nc == 0:
            transform_A = get_transform(self.opt, params)
            A_tensor = transform_A(A.convert('RGB'))
        else:
            transform_A = get_transform(self.opt, params, method=Image.NEAREST, normalize=False)
            A_tensor = transform_A(A) * 255.0

        M_path = self.M_paths[index]
        M = Image.open(M_path).convert('L')
        params = get_params(self.opt, M.size)
        if self.opt.label_nc == 0:
            transform_M = get_transform(self.opt, params)
            M_tensor = transform_M(M.convert('RGB'))
            detail_tensor = transform_M(detail.convert('RGB'))
        else:
            transform_M = get_transform(self.opt, params, method=Image.NEAREST, normalize=False)
            M_tensor = transform_M(M) * 255.0

        B_tensor = inst_tensor = feat_tensor = 0
        if self.opt.isTrain or self.opt.use_encoded_image:
            B_path = self.B_paths[index]   
            B = Image.open(B_path).convert('RGB')
            transform_B = get_transform(self.opt, params)      
            B_tensor = transform_B(B)

        ### if using instance maps        
        if not self.opt.no_instance:
            inst_path = self.inst_paths[index]
            inst = Image.open(inst_path)
            inst_tensor = transform_A(inst)

            if self.opt.load_features:
                feat_path = self.feat_paths[index]            
                feat = Image.open(feat_path).convert('RGB')
                norm = normalize()
                feat_tensor = norm(transform_A(feat))                            

        input_dict = {'label': A_tensor, 'inst': inst_tensor, 'image': B_tensor, 'mask':M_tensor,'detail':detail_tensor,
                      'feat': feat_tensor, 'path': A_path}

        return input_dict

    def __len__(self):
        return len(self.A_paths) // self.opt.batchSize * self.opt.batchSize

    def name(self):
        return 'AlignedDataset'